import pandas as pd
import joblib
import numpy as np


def predict_with_lightgbm(models_file, data_file, forecast_days=14):
    """使用训练好的LightGBM模型预测未来销量"""
    # 加载模型
    models = joblib.load(models_file)

    # 加载历史数据
    df = pd.read_excel(data_file)

    # 存储预测结果
    predictions = []

    # 对每个款号进行预测
    for spu, model in models.items():
        # 获取该款号的最新数据
        spu_data = df[df['spu_code'] == spu]
        last_record = spu_data.iloc[-1]
        last_day = last_record['days_since_launch']

        # 预测未来14天
        for day in range(1, forecast_days + 1):
            # 准备预测特征
            features = pd.DataFrame({
                'sum_cart_total': [last_record['sum_cart_total']],
                'new_sale_qty': [last_record['new_sale_qty']],
                'days_since_launch': [last_day + day]
            })

            # 预测销量
            pred_sale = model.predict(features)[0]

            # 确保非负
            pred_sale = max(0, pred_sale)

            predictions.append({
                'spu_code': spu,
                'days_since_launch': last_day + day,
                'predicted_sale_qty': pred_sale
            })

    # 转换为DataFrame
    return pd.DataFrame(predictions)


if __name__ == "__main__":
    # 执行预测
    future_sales = predict_with_lightgbm(
        models_file='lightgbm_models.pkl',
        data_file='sales_data.xlsx'
    )

    # 保存结果
    future_sales.to_excel("enhanced_sales_data.xlsx", index=False)
    print("预测结果已保存为 enhanced_sales_data.xlsx")
    print("\n预测结果示例:")
    print(future_sales.head(14))